#!/bin/bash

CURRENT_DIR=$(cd "$(dirname "$0")"; pwd)

declare -A index_model_map=(
    ["1"]="Qwen/Qwen1.5-1.8B"
    ["2"]="Qwen/Qwen2.5-3B-Instruct"
    ["3"]="Qwen/Qwen3-0.6B"
    ["4"]="Qwen/Qwen3-8B"
    ["5"]="JunHowie/Qwen3-8B-GPTQ-Int4"
    ["6"]="AngelSlim/Qwen3-8b_int4_awq"
    ["7"]="soob3123/amoral-qwen3-8B-Q4_K_M-GGUF"

)

# ask user to select model
read  -t 6 -p "[Q&A] Please select the index of the model to download, [-1] - help, [0] - all, [1] - Qwen1_5-1_8B, [2] - Qwen2.5-3b-instruct : " model_index
#echo "model: ${index_model_map[$model_index]}"

# display help information
if [ -z "$model_index" ] || [ "$model_index" == "-1" ]; then
    echo -e "\nAvailable models :"
    for index in "${!index_model_map[@]}"; do
        echo "[$index] - ${index_model_map[$index]}"
    done
    exit 0
fi

# check if 'root' user
if [ "$(id -u)" -eq 0 ]; then
    apt install -y python3-dev
    apt install -y python3-pip
else
    sudo apt install -y python3-dev
    sudo apt install -y python3-pip
fi
pip3 install modelscope -i https://pypi.tuna.tsinghua.edu.cn/simple/

case $model_index in
    0)
        for model in "${index_model_map[@]}"; do
            # check if directory exists
            if [ -d "${CURRENT_DIR}/../model/${model}" ]; then
                rm -rf ${CURRENT_DIR}/../model/${model}
            fi

            # download model
            echo "[INFO] Downloading model: $model"
            modelscope download --model ${model} --local_dir "${CURRENT_DIR}/../model/${model}"
        done
        ;;
    1|2|3|4|5|6|7)
        model=${index_model_map[$model_index]}
        # check if directory exists
        if [ -d "${CURRENT_DIR}/../model/${model}" ]; then
            rm -rf ${CURRENT_DIR}/../model/${model}
        fi

        # download model
        echo "[INFO] Downloading model: $model"
        modelscope download --model ${model} --local_dir "${CURRENT_DIR}/../model/${model}"
        ;;
    *)
        echo "[ERROR] Invalid input : ${model_index}. Please enter a valid index."
        exit 1
        ;;
esac